
#include "aa_neighborhood.h"
#include "aa_util.h"
#include "../pdb_index/aa_table.h"
#include "permutations.h"

#include "../query_handler/query_handler.h"
#include "translationRotation.h"

#include <math.h>
#include <values.h>
#include <assert.h>

//int coordinates3_calculate(struct coordinates3* c, const char pdbName[4], const struct store * s, const off_t position);

void aa_neighborhood_find_neighbors(struct mTree * tree, const struct store * s, struct atom * center,  double radius, 
				    struct aa_neighborhood * neigh, short count) ;



int coordinates3_calculate(struct xyz_short* c, const char pdbName[4], const struct store * s, const off_t position) {
  off_t i = find_in_aa_by_name(s, position, pdbName);

  if(i == -1) {
     return 0;
  } else {
    struct atom * a = store_read(s, i);
    
    c->x = a->x/10;
    c->y = a->y/10;
    c->z = a->z/10;
    
    store_release(s,i);
    return 1;
  }
}

/*double coordinates3_distance(const struct coordinates3 * c1, const struct coordinates3 * c2) {
  double x = c1->x - c2->x;
  double y = c1->y - c2->y;
  double z = c1->z - c2->z;

  x*=x;
  y*=y;
  z*=z;

  return sqrt(x+y+z);  
  }*/

#ifdef NEIGHBORS
int aa_neighborhood_permutation_array[NEIGHBORS];

double min_pairwise_distance(const struct xyz_short n1[], const type n1_t[],
			     const struct xyz_short n2[], const type n2_t[],
			     short count) {
  short i = 0;
  double min_dist, dist;

  dist = 0;
  for(i = 0; i < count; i++) {
    aa_neighborhood_permutation_array[i] = i;

    dist += xyz_short_distance(&(n1[i]), &(n2[i]));

    //    dist += aa_type_distance(n1_t[i], n2_t[i]);

  }

  min_dist = dist;


  while(permutation_next(aa_neighborhood_permutation_array, count)) {
    dist = 0.0;
    for(i = 0; i < count; i++) {
      dist += xyz_short_distance(&(n1[i]), &(n2[aa_neighborhood_permutation_array[i]]));
      //      dist += aa_type_distance(n1_t[i], n2_t[aa_neighborhood_permutation_array[i]]);
    }
    if(dist < min_dist) {
      min_dist = dist;
    }
  }
	
  return min_dist;
}
#endif // NEIGHBORS
double aa_neigh_dist(const struct aa_neighborhood* n1, const struct aa_neighborhood* n2) {
  double ret = 0.0;
#ifdef CENTER_TYPE
  ret += CENTER_TYPE_WEIGHT      * aa_type_distance(n1->center_type, n2->center_type);
  assert(ret < 1e90);
#endif
#ifdef NEIGHBORS
  ret += NEIGHBORS_WEIGHT        * min_pairwise_distance(n1->neighbors, n1->neighbor_types,
							 n2->neighbors, n2->neighbor_types, NEIGHBORS);
  assert(ret < 1e90);
#endif
#ifdef MASS_CENTER
  ret += MASS_CENTER_WEIGHT      * xyz_short_distance(&(n1->mass_center), &(n2->mass_center));
  assert(ret < 1e90);
#endif
#ifdef FUNCTIONAL_GROUP
  ret += FUNCTIONAL_GROUP_WEIGHT * xyz_short_distance(&(n1->functional_group), &(n2->functional_group));
  assert(ret < 1e90);
#endif
#ifdef FIELDS
  ret += FIELDS_WEIGHT           * fields_distance(n1->fields, n1->field_types, n2->fields, n2->field_types, FIELDS);
  assert(ret < 1e90);
#endif
#ifdef BONDS
  ret += BONDS_WEIGHT            * bonds_distance(n1->bonds, n1->bond_types, n2->bonds, n2->bond_types, BONDS);
  assert(ret < 1e90);
#endif
  
  assert(ret < 1e90);

  return ret;
}

void aa_neigh_init(struct aa_neighborhood * neigh, const char pdbName[4], const struct store * s, struct mTree * atom_tree, off_t position) {
  double trans_vect_array[3];
  double rot_matrix_array[3 * 3];
  short i;
  gsl_vector_view translation_vector = gsl_vector_view_array(trans_vect_array, 3);
  gsl_matrix_view rotation_matrix = gsl_matrix_view_array(rot_matrix_array, 3, 3);

  struct atom * center = store_read(s, position);
  struct xyz_int tmp_xyz;
  off_t start, stop;

  struct atom mass_center;

#ifdef FIELDS
  double weights[NEIGHBORS];
#endif


#ifdef REFERENCE_FRAME_N_CA_C
  struct xyz_short ca ;
  struct xyz_short n  ;
  struct xyz_short c  ;
  if(
     coordinates3_calculate(&ca, " CA ", s, position) &&
     coordinates3_calculate(&n , " N  ", s, position) &&
     coordinates3_calculate(&c , " C  ", s, position) ) {
    calculate_reference_frame(&ca, &n, &c, &(translation_vector.vector), &(rotation_matrix.matrix));
  } else {
    assert(0);
  }
#endif
#ifdef REFERENCE_FRAME_CA_CB_C
  struct xyz_short ca ;
  struct xyz_short cb ;
  struct xyz_short c  ;
  if( 
     coordinates3_calculate(&ca, " CA ", s, position) &&
     coordinates3_calculate(&n , " CB ", s, position) &&
     coordinates3_calculate(&c , " C  ", s, position) ) {
    calculate_reference_frame(&ca, &cb, &c, &(translation_vector.vector), &(rotation_matrix.matrix));
  } else {
    assert(0);
  }
#endif

  neigh->aa_index    = center->residue;
  strncpy(neigh->pdb_name, pdbName, 4);

#ifdef CENTER_TYPE

  neigh->center_type = center->residue_type;

#endif

  
#ifdef NEIGHBORS

  mass_center.x = -trans_vect_array[0] * 10;
  mass_center.y = -trans_vect_array[1] * 10;
  mass_center.z = -trans_vect_array[2] * 10;
  mass_center.residue = center->residue;

  aa_neighborhood_find_neighbors(atom_tree, s, &(mass_center), 30 * 1000, neigh, NEIGHBORS);

  for(i = 0; i < NEIGHBORS; i++) {
    translate_and_rotate(&(neigh->neighbors[i]), &(translation_vector.vector), &(rotation_matrix.matrix));
    translate_and_rotate(&(neigh->neighbor_dirs[i]), &(translation_vector.vector), &(rotation_matrix.matrix));
  }

#endif
#ifdef MASS_CENTER

  find_aa_range(s, position, &start, &stop);

  calculate_aa_mass_center(s, start, stop, &(tmp_xyz));
  neigh->mass_center.x = tmp_xyz.x / 10;
  neigh->mass_center.y = tmp_xyz.y / 10; 
  neigh->mass_center.z = tmp_xyz.z / 10;

  translate_and_rotate(&(neigh->mass_center), &(translation_vector.vector), &(rotation_matrix.matrix));

#endif
#ifdef FUNCTIONAL_GROUP

#ifndef MASS_CENTER 
  find_aa_range(s, position, &start, &stop);
#endif

  if (1 == calculate_aa_functional_group(s, start, stop, &(tmp_xyz))) {
    neigh->functional_group.x = tmp_xyz.x /10;
    neigh->functional_group.y = tmp_xyz.y /10;
    neigh->functional_group.z = tmp_xyz.z /10;
    translate_and_rotate(&(neigh->functional_group), &(translation_vector.vector), &(rotation_matrix.matrix));
  } else {
    neigh->functional_group.x = neigh->mass_center.x;
    neigh->functional_group.y = neigh->mass_center.y;
    neigh->functional_group.z = neigh->mass_center.z;
  }


#endif
#ifdef FIELDS

  for(i = 0; i < NEIGHBORS; i++) {
    weights[i] = aa_type_hydrophobicity(neigh->neighbor_types[i]);
  }

  calculate_hydrophobicity_field(center, neigh->neighbors, weights, NEIGHBORS, neigh->fields[0], &(rotation_matrix.matrix));
  //  calculate_charge_field(center, charged_coordinates, charged_weights, charged_count, neigh->fields[1], &(rotation_matrix.matrix));

#endif
#ifdef BONDS

  find_bonds(s, bs, position, neigh->bonds, BONDS); // bs is a bondStore for this structure...

  for(i = 0; i < BONDS; i++) {
    translate_and_rotate(&(neigh->bonds[i]), &(translation_vector.vector), &(rotation_matrix.matrix));
  }
#endif

  store_release(s, position);

}


void aa_neighborhood_find_neighbors(struct mTree * tree, const struct store * s, struct atom * center,  double radius, 
				    struct aa_neighborhood * neigh, short count) { 

  struct knn_state * knn_state = knn_state_alloc(count, center, (double(*)(const void *, const void *))atom_distance, sizeof(struct atom));
  struct filter_state filter_state;
  short c;

  knn_state_reset(knn_state);

  filter_state.count        = 3;
  filter_state.functions    = malloc(sizeof(void *(*)(void *, void *)) * 3);
  filter_state.states       = malloc(sizeof(void *) * 3);
  //  filter_state.functions[0] = (void *(*)(void*, void*))filter_atom_by_name;
  filter_state.functions[0] = (void *(*)(void*, void*))filter_atom_by_name;
  filter_state.states[0]    = " CA ";
  filter_state.functions[1] = (void *(*)(void*, void*))filter_atom_by_residue_number;
  filter_state.states[1]    = (void *)center->residue;
  filter_state.functions[2] = (void *(*)(void*, void*))knn;
  filter_state.states[2]    = knn_state;


  mTree_rangeQuery(tree, center, radius * 100.0, (void(*)(void*,void*))filterAnd, &filter_state);

  for(c = 0; c < count; c++) {

    if(knn_state->results_distance[c] != MAXDOUBLE) {

      struct atom * neighbor_atom = &((struct atom*)knn_state->results)[c];
      off_t start, stop;
      struct xyz_int tmp_xyz;

      find_aa_range(s, neighbor_atom->serialNumber, &start, &stop);

      calculate_aa_mass_center(s, start, stop, &(tmp_xyz));

      neigh->neighbors[c].x = tmp_xyz.x / 10;
      neigh->neighbors[c].y = tmp_xyz.y / 10;
      neigh->neighbors[c].z = tmp_xyz.z / 10;

      // Rotation handler by caller.

      if(1 == calculate_aa_functional_group(s, start, stop, &(tmp_xyz))) {
	neigh->neighbor_dirs[c].x = tmp_xyz.x / 10 ;
	neigh->neighbor_dirs[c].y = tmp_xyz.y / 10 ;
	neigh->neighbor_dirs[c].z = tmp_xyz.z / 10 ;
      } else {
	neigh->neighbor_dirs[c].x = neigh->neighbors[c].x;
	neigh->neighbor_dirs[c].y = neigh->neighbors[c].y;
	neigh->neighbor_dirs[c].z = neigh->neighbors[c].z;
      }
      neigh->neighbor_types[c]   = neighbor_atom->residue_type;
      neigh->neighbor_indices[c] = neighbor_atom->residue; 

    }

  }
  knn_state_free(knn_state);
  free(filter_state.functions);
  free(filter_state.states);


}

void calculate_hydrophobicity_field(const struct xyz_short * center, const struct xyz_short coords[], const double weights[],
				    aa_index count, gsl_matrix * rotation_matrix, struct xyz_short * out) {

  aa_index i; 
  float s_x = 0, s_y = 0, s_z = 0, s_w = 0;
  
  for(i = 0; i < count; i++) {
    s_x += (coords[i].x - center->x) * weights[i];
    s_y += (coords[i].y - center->y) * weights[i];
    s_z += (coords[i].z - center->z) * weights[i];
    s_w += weights[i];
  }

  s_x /= s_w;
  s_y /= s_w;
  s_z /= s_w;
  
  out->x = (short) s_x;
  out->y = (short) s_y;
  out->z = (short) s_z;

  rotate(out, rotation_matrix);

}

/**
   center should be in it's original position, as should coords[].  weights[] will generally be the electric chare for each amino acid type in coords.
*/
void calculate_charge_field(const struct xyz_short * center, const struct xyz_short coords[], const float weights[], aa_index count, gsl_matrix * rotation_matrix, struct xyz_short * out) {

  aa_index i;
  float sum_forces_x = 0, sum_forces_y = 0, sum_forces_z = 0;

  for(i = 0; i < count; i++) {
    float d;
    float uv_x, uv_y, uv_z;


    uv_x = center->x - coords[i].x;  
    uv_y = center->y - coords[i].y;
    uv_z = center->z - coords[i].z;

    /* uv_* aren't unit vectors yet--Use them to calculate d... */

    d = (float)sqrt((uv_x * uv_x) + (uv_y * uv_y) + (uv_z * uv_z)); 

    uv_x /= d;
    uv_y /= d;
    uv_z /= d;

    sum_forces_x += (weights[i] * uv_x / (d*d));
    sum_forces_y += (weights[i] * uv_y / (d*d));
    sum_forces_z += (weights[i] * uv_z / (d*d));

  }

  /* sum_forces_* are now centered about the origin of the reference frame, but they still need to be rotated... */
  
  out->x = (short) sum_forces_x;
  out->y = (short) sum_forces_y;
  out->z = (short) sum_forces_z;

  rotate(out, rotation_matrix);

}

int aa_neigh_can_calc_neighborhood(struct store * s, off_t i) {
#ifdef REFERENCE_FRAME_N_CA_C
  return (find_in_aa_by_name(s, i, " CA ") != -1) && 
    (find_in_aa_by_name(s, i, " N  ") != -1) && 
    (find_in_aa_by_name(s, i, " C  ") != -1);
#endif
#ifdef REFERENCE_FRAME_CA_CB_C
  return (find_in_aa_by_name(s, i, " CA ") != -1) && 
    (find_in_aa_by_name(s, i, " CB ") != -1) && 
    (find_in_aa_by_name(s, i, " C  ") != -1);
#endif
}
